package org.rx;

import java.io.FileNotFoundException;
import java.io.InputStreamReader;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Comparator;
import java.util.HashSet;

import org.antlr.runtime.ANTLRFileStream;
import org.antlr.runtime.ANTLRReaderStream;
import org.antlr.runtime.ANTLRStringStream;
import org.antlr.runtime.CommonTokenStream;
import org.antlr.runtime.RecognitionException;
import org.antlr.runtime.TokenRewriteStream;
import org.antlr.runtime.tree.CommonTree;
import org.antlr.runtime.Token;
import org.rx.SXPNumbering.Abstract;
import org.rx.SXPNumbering.Concrete;
import org.rx.SXPNumbering.Node;

public class Dispatcher {
	static public final String inputExt[] = {".txt", ".dispatch"};
	static public String output = System.getProperty("extension", ".inc");
	
	CommonTree method;
	MethodEntry entries[];
	
	Poset<MethodEntry> poset = new Poset<MethodEntry>();
	boolean inline, swap;
	DispatcherType dispatchType = DispatcherType.AUTO;
	String template = "";
		
	Dispatcher(CommonTree m){
		method = m;
		int nb = m.getChildCount();
		ArrayList<MethodEntry> ent = new ArrayList<MethodEntry>();
		for(int i = 2; i < nb; i++){
			CommonTree def = (CommonTree)m.getChild(i);
			String t = def.getText();
			if(def.getType() == DispatcherParser.DEF) {
				MethodEntry me = new MethodEntry(def);
				if(ent.contains(me))
					System.err.println("Duplicate entry found for: "+me);
				ent.add(me);
			} else if(def.getType() == DispatcherParser.TEMPLATENAME){
				template = def.getText();
			} else if(t.equalsIgnoreCase("INLINE")){
				inline = true;
			} else if(t.equalsIgnoreCase("SWAP")){
				swap = true;
			} else if(t.equalsIgnoreCase("SWITCH")){
				dispatchType = DispatcherType.SWITCH;
			} else if(t.equalsIgnoreCase("BTD")){
				dispatchType = DispatcherType.BTD;
			} else if(t.equalsIgnoreCase("TABLE")){
				dispatchType = DispatcherType.TABLE;
			} else {
				System.err.println("Ignoring unsuported modifier: "+t);
			}
		}
		entries = new MethodEntry[ent.size()];
		ent.toArray(entries);
	}
	
	boolean isVoid(){
		return rtype().equals("void");
	}
	String name(){
		return method.getToken().getText();
	}
	
	String rtype(){
		return typeToString((CommonTree)method.getChild(0));
	}
	int nbArgs(){
		return sigNode().getChildCount();
	}
	CommonTree sigNode(){
		return (CommonTree)method.getChild(1);
	}
	String signature(){
		return signatureToString(sigNode());
	}
	String signatureNamed(String prefix){	
		return signatureToStringNamed(sigNode(), prefix);
	}
	String functionSignature(String f, String s){
		return rtype() + " " + name() +"("+f+", "+s + signature()+")";
	}

	String functionSignatureNamed(String f, String s, String l, String r, String p){
		return rtype() + " " + name() +"("+f+" "+l+", "+s + " "+r+signatureNamed(p)+")";
	}
	
	String ptrSignature(String ptr, String f, String s){
		return rtype() + " (*" + ptr +") ("+f+", "+s + signature()+")";
	}
	
	DispatcherType dispatcherType(){
		// TODO select the most appropriate
		if(dispatchType == DispatcherType.AUTO)
			return DispatcherType.SWITCH;
		return dispatchType;
	}
	void generate(PrintStream out, String l, String r, String p){
		Collection<Node> set = new HashSet<Node>();
		for(MethodEntry me: entries){
			set.add(me.left);
			set.add(me.right);
		}
		
		Node root = Hierarchy.hierarchy.supertype_of_all(set);
		Shift s = new Shift(root);
		if(template != null)
			out.println(template);
		if(inline)
			out.print("inline ");
		out.print(functionSignatureNamed("Any *", "Any *", l, r, p));
		out.println(" {");
			generatePrologue(out, s, l, r, p);
			generateDispatcher(out, s, l, r, p);
		out.println("}");
	}

	
	void generatePrologue(PrintStream out, Shift s, String l, String r, String p){
		out.println("\tllvm_gcroot("+l+", 0);");
		out.println("\tllvm_gcroot("+r+", 0);");
		
		out.println("\tint ltype = "+l+"->getType("+l+");");
		out.println("\tint rtype = "+r+"->getType("+r+");");
		
		if(swap){
			out.println("\tif(ltype < rtype){");
			out.println("\t\tswap("+l+", "+r+");");
			out.println("\t\tswap(ltype, rtype);");
			out.println("\t}");
		}
		if(s.root != Hierarchy.hierarchy){
			out.println("\tif(!"+s.root.isA("ltype")+") " +dispatchError()+callArgs(l, Hierarchy.any, r, Hierarchy.any, p)+";");
			out.println("\tif(!"+s.root.isA("rtype")+") " +dispatchError()+callArgs(l, Hierarchy.any, r, Hierarchy.any, p)+";");
		}
	}
	
	void generateDispatcher(PrintStream out, Shift s, String l, String r, String p){
		out.println("\tint pos = " + s.generateCCode("ltype", "rtype") + ";");
		switch(dispatcherType()){
		case TABLE:
			generateTable(out, s, l, r, p); // FIXME
			break;
		case SWITCH:
			generateSwitch(out, s, l, r, p);
			break;
		case BTD:
		default:
			out.println("\tnyi();"); // TODO
			break;
		}
	}

	String dispatchError(){
		return "dispatchError";
	}
	void generateTable(PrintStream out, Shift s, String l, String r, String p){
		int nbComb = (2 << s.len - 1);
		String table[] = new String[nbComb * nbComb];
		Arrays.fill(table, dispatchError());
		for(MethodEntry m: entries)
			table[s.computeShift(m.left.id, m.right.id)] = m.target;	
		String type = s.root.name+"*";
		out.println("\tstatic const "+ptrSignature("table", type, type)+"[] = {"+arrayToString(table, ", ")+"}");
		if(!isVoid())
			out.print("\treturn ");
		else
			out.print("\t");
		out.println("table[pos]"+callArgs(l, Hierarchy.any, r, Hierarchy.any, p) + ";");
	}
	
	void generateSwitch(PrintStream out, Shift s, String l, String r, String p){
		boolean isVoid = isVoid();
		out.println("\tswitch(pos){");
		for(MethodEntry m: entries){
			generateCaseFor(out, s, l, m.left, r, m.right, p, isVoid, m.target, false);
			if(!swap && getEntry(m.right, m.left) == null)
				generateCaseFor(out, s, l, m.left, r, m.right, p, isVoid, m.target, true);
		}
		out.println("\t\tdefault: "+dispatchError()+callArgs(l, Hierarchy.any, r, Hierarchy.any, p)+";");
		if(!isVoid) out.println("\t\t\treturn 0;"); // FIXME
		out.println("\t};");
	}
	
	void generateCaseFor(PrintStream out, Shift s, String l, Node left, String r, Node right, String p, boolean isVoid, String target, boolean rev){
		String call;
		if(rev){
			call = target+callArgs(r, left, l, right, p)+";";
			out.println("\t\tcase "+s.computeShift(left.id, right.id)+":");
		}else{
			call = target+callArgs(l, left, r, right, p)+";";
			out.println("\t\tcase "+s.computeShift(right.id, left.id)+":");
		}
		
		if(isVoid){
			if(rev)
				out.println("\t\t\t"+call);
			else
				out.println("\t\t\t"+call);
			out.println("\t\t\treturn;");
		}else
			if(rev)
				out.println("\t\t\treturn "+call);
			else
				out.println("\t\t\treturn "+call);
	}
	
	String callArgs(String l, Node ltype, String r, Node rtype, String p){
		int nbargs = nbArgs();
		StringBuffer buff = new StringBuffer("(static_cast<").append(ltype.name).append("*>(").append(l).
							append("), static_cast<").append(rtype.name).append("*>(").append(r).append(")");
		for(int i = 0; i < nbargs; i++)
			buff.append(", ").append(p).append(i);
		return buff.append(")").toString();
	}
	
	private void makeLatice() {
		Arrays.sort(entries, signatureComparator);
		ArrayList<MethodEntry> supers = new ArrayList<MethodEntry>();
		for(int i = 0; i < entries.length; i ++){
			for(int j = 0; j < i; j ++){
				if(signatureComparator.compare(entries[j], entries[i]) < 0){
					supers.add(entries[j]);
				}
			}
			poset.add(entries[i], supers);
			supers.clear();
		}
		poset.validate();
		//System.out.println(poset.toDot());
	}
	
	static Comparator<MethodEntry> signatureComparator = new Comparator<MethodEntry>(){	
		public int compare(MethodEntry one, MethodEntry other){
			if(one.left.id < other.left.id && one.right.id < other.right.id
					|| one.left.id < other.left.id && one.right.id == other.right.id
					|| one.left.id == other.left.id && one.right.id < other.right.id)
				return -1;
			if(one.left.id > other.left.id && one.right.id > other.right.id
					|| one.left.id > other.left.id && one.right.id == other.right.id
					|| one.left.id == other.left.id && one.right.id > other.right.id)
				return 1;

			return 0;
		}
	};
	
	public static void main(String[] args) throws ClassNotFoundException {
		Class.forName("org.rx.SXPNumbering");
		// not prety but ensure to numbering BEFORE generating eveything
		PrintStream out = System.out;
		for (String arg: args) {
			ANTLRStringStream reader;
			try {
				if(arg.equals("-")){
					reader = new ANTLRReaderStream(new InputStreamReader(System.in));
				} else {
					reader = new ANTLRFileStream(arg);
					out = outputStream(arg);
				}
				process_file(reader, out, arg);
			} catch (Exception e) {
				System.err.println("Parse error Exception:"+e.getMessage());
			}
		}
	}

	private static void process_file(ANTLRStringStream reader, PrintStream out, String fname) throws RecognitionException {
		DispatcherLexer lexer = new DispatcherLexer(reader);
		// CommonTokenStream tokens = new CommonTokenStream(lexer);
		CommonTokenStream tokens = new TokenRewriteStream(lexer);
		DispatcherParser parser = new DispatcherParser(tokens);
		DispatcherParser.start_return result = parser.start();
		CommonTree tree;
		if(result == null || (tree = (CommonTree)result.getTree()) == null){
			return;
		}
		out.println("// Generated from "+fname);
		for(Object t: tree.getChildren()){
			try {
				Dispatcher d = new Dispatcher((CommonTree)t);
				d.makeLatice();
				d.generate(out, "left", "right", "param");
			} catch (Exception e) {
				e.printStackTrace();
			}
		}
	}

	static PrintStream outputStream(String name) throws FileNotFoundException{
		if(System.getProperty("stdoutput") != null) // TODO true options
			return System.out;
		
		for(String ext: inputExt)
			if(name.endsWith(ext))
				name = name.substring(0, name.length() - ext.length());

		return new PrintStream(name+".inc");
	}
	
	private static String typeToString(CommonTree tree){
		return tree.getText();
	}
	
	private static String signatureToString(CommonTree tree){
		if(tree == null || tree.getChildCount() == 0)
			return "";
		StringBuffer buff = new StringBuffer();
		for(Object c: tree.getChildren())
			buff.append(", ").append(typeToString((CommonTree)c));
		return buff.toString();
	}
	
	private static String signatureToStringNamed(CommonTree tree, String prefix){
		if(tree == null || tree.getChildCount() == 0)
			return "";
		StringBuffer buff = new StringBuffer();
		int i = 0;
		for(Object c: tree.getChildren())
			buff.append(", ").append(typeToString((CommonTree)c)).append(" ").append(prefix).append(i++);
		return buff.toString();
	}
	
	
	class Shift{
		ArrayList<Concrete> set;
		int mask, shift, len;
		Node root;
		
		Shift(Node root2){
			root = root2;
			set = new ArrayList<Concrete>();
			root.getLeaves(set);
			int ids[] = new int[set.size()];
			for(int i = 0; i < ids.length; i++)
				ids[i] = set.get(i).id;
			mask = Hierarchy.discriminantsBits(ids);
			shift = Hierarchy.rightMost1Bit(mask) - 1;
			len = Integer.bitCount(mask);
		}
		String generateCCode(String l, String r){
			int oshift = shift - len;
			if(oshift > 0)
				return "(("+l+"&"+mask+")>>"+shift+") | (("+r+"&"+mask+") >> ("+(oshift)+"))";
			else
				return "(("+l+"&"+mask+")>>"+shift+") | (("+r+"&"+mask+") << ("+(-oshift)+"))";

		}
		
		int computeShift(int l, int r){
			int oshift = shift - len;
			if(oshift > 0)
				return ((l& mask)>>shift) | ((r&mask) >> (oshift));
			return ((l& mask)>>shift) | ((r&mask) << (-oshift));
		}
	}
	
	MethodEntry getEntry(Node l, Node r){
		for(MethodEntry me: entries)
			if(l == me.left && r == me.right)
				return me;
		return null;
	}
	class MethodEntry {
		Node left;
		Node right;
		String target;
		
		MethodEntry(CommonTree def){
			target = def.getToken().getText();
			String l = typeToString((CommonTree)def.getChild(0));
			String r = typeToString((CommonTree)def.getChild(1));
			left = Hierarchy.toNode(l);
			right = Hierarchy.toNode(r);
			if(left == null || right == null)
				throw new RuntimeException("One (or more) of this type doesn't belong to the hierarchy: "+l+ " "+r);
			if(left instanceof Abstract || right instanceof Abstract)
				throw new RuntimeException("Dispatcher currently only works for concrete types: "+left+ " "+right);
		}
		
		public String toString(){
			return functionSignature(left.name, right.name);
		}
		
		public boolean equals(Object o){
			if(o instanceof MethodEntry){
				MethodEntry me = (MethodEntry)o;
				return left == me.left && right == me.right;
			}
			return false;
		}
	}
		
	public static void printTree(CommonTree tree, PrintStream out, int level) {
		try {
			Token token = tree.getToken();
			System.out.println();
			for(int i = level; i > 0; i--)
				System.out.print('\t');
			if(tree.getChildCount() > 0){
				if(token == null)
					System.out.print("(token:null");
				else
					out.print("("+token.getText()+":"+token.getType());
				if(tree.getChildren() != null)
					for(Object c: tree.getChildren()){
						printTree((CommonTree)c, out, level + 1);
					}
				out.print(")");
			} else
				out.print(token.getText()+":"+token.getType());				
		} catch (Exception e) {
			e.printStackTrace();
		}
		if(level == 0)
			out.println();
	}
	
	public static String arrayToString(String array[], String sep){
		StringBuffer buff = new StringBuffer(array[0].toString());
		for(int i = 1; i < array.length; i++)
			buff.append(sep).append(array[i].toString());
		return buff.toString();
	}
	
	public enum DispatcherType {
	    AUTO, TABLE, SWITCH, BTD
	}
}
